# Copyright 2023 ClickMab Biotech Suzhou Co., LTD
# 串联 INSPIRED 相关的各个模块，完成对抗体 binning 的完整分析流程

import os
import shutil
import argparse
import subprocess
import logging
import json


logger = logging.getLogger(__name__)
logger.setLevel(logging.INFO)
formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
handler = logging.StreamHandler()
handler.setFormatter(formatter)
logger.addHandler(handler)

# 获取流程涉及的镜像
current_folder = os.path.dirname(os.path.realpath(__file__))
with open(f'{current_folder}/images.json', encoding='utf-8') as fi:
    IMAGES = json.load(fi)


def is_docker_image(image_name):
    """检查镜像是否存在"""
    command = f'docker inspect {image_name}'
    result = subprocess.run(command, shell=True, check=False, stdout=subprocess.PIPE, stderr=subprocess.PIPE)

    if result.returncode == 0:
        return True
    else:
        return False


def pipeline(ag_pdb_fp: str,
             ab_csv_fp: str,
             gpu='0',
             k=[5, 10, 15],
             save_fd='./inspired-result',
             tmp_fd='./.inspired_tmp_data'):
    """
    构建抗体 binning 分析的流程。

    Parameters
    ----------
    ag_pdb_fp : str
        抗原结构 pdb 文件路径
    ab_csv_fp : str
        抗体信息 csv 文件路径，包含三列 name, VH 和 VL
    gpu : str
        gpu device id
    k : list[int]
        挑选用于 binning 的抗体数量
    save_fd : str
        保存结果文件夹路径
    """
    logger.info('Preparing environment..')
    # 检查依赖的镜像是否都存在
    for image_name in IMAGES.values():
        if not is_docker_image(image_name):
            raise RuntimeError(f'Docker image {image_name} not found!')
    # 中间分析结果的文件夹
    tmp_fd = os.path.realpath(tmp_fd)
    # if os.path.isdir(tmp_fd):
    #     shutil.rmtree(tmp_fd)
    os.makedirs(tmp_fd,exist_ok=True)

    # 0. Check data
    logger.info('Check data..')
    check_data_fd = f'{tmp_fd}/0-check_data'
    os.makedirs(check_data_fd,exist_ok=True)
    # 准备输入数据
    os.makedirs(f'{check_data_fd}/input',exist_ok=True)
    shutil.copy(ag_pdb_fp, f'{check_data_fd}/input/ag.pdb')
    shutil.copy(ab_csv_fp, f'{check_data_fd}/input/ab.csv')
    # 准备输出文件夹
    os.makedirs(f'{check_data_fd}/output',exist_ok=True)
    # 调用命令
    cmd = f"""
docker run --rm -v {check_data_fd}/input:/job/data \\
-v {check_data_fd}/output:/job/data/output \\
{IMAGES.get("check_data")} \\
    -d /job/data \\
    -o /job/data/output
""".strip()

    os.system(cmd)

    # 1.predict Ab struct
    logger.info('Predict antibody structure..')
    ab_struct_predict_fd = f'{tmp_fd}/1-ab_struct_predict_and_cutPatch'
    os.makedirs(ab_struct_predict_fd, exist_ok=True)
    # 准备输入数据
    ab_struct_predict_input = f'{ab_struct_predict_fd}/input'
    os.makedirs(ab_struct_predict_input,exist_ok=True)
    shutil.copy(f"{check_data_fd}/output/ag.pdb", f'{ab_struct_predict_input}/ag.pdb')
    shutil.copy(f"{check_data_fd}/output/ag.csv", f'{ab_struct_predict_input}/ag.csv')
    # 准备输出文件夹
    ab_struct_predict_output = f'{ab_struct_predict_fd}/output'
    os.makedirs(ab_struct_predict_output,exist_ok=True)
    if gpu == '-1':
        cmd = f"""
docker run --rm \\
-v {ab_struct_predict_input}:/job/Testdata \\
-v {ab_struct_predict_output}:/job/data/patchData \\
-e CUDA_VISIBLE_DEVICES={gpu} \\
{IMAGES.get("ab_struct_predict")} \\
    -d /job/Testdata \\
    -t immunebuilder
""".strip()
    else:
        cmd = f"""
docker run --rm --gpus all \\
-v {ab_struct_predict_input}:/job/Testdata \\
-v {ab_struct_predict_output}:/job/data/patchData \\
-e CUDA_VISIBLE_DEVICES={gpu} \\
{IMAGES.get("ab_struct_predict")} \\
    -d /job/Testdata \\
    -t immunebuilder
""".strip()
    os.system(cmd)

    # # 2. Cut patches
    logger.info('Cut patches..')
    cut_patch_fd = f'{tmp_fd}/2-cut_patch' # 暂时未用
    os.makedirs(cut_patch_fd,exist_ok=True)
    # 调用命令
    cmd = f"""
docker run --rm \\
-v {ab_struct_predict_fd}/output:/job/data/patchData \\
-v {ab_struct_predict_fd}/input:/job/data/pdb_csv \\
{IMAGES['cut_patches']} \\
    -d /job/data/pdb_csv \\
    -a /job/data/patchData/*/Igfold
""".strip()
    os.system(cmd)
    
    # 3. Predict Ag-patch ab scores
    logger.info('Predict scores of Antigen-patches of antibodies.')
    predict_ab_fd = f'{tmp_fd}/3-predict_ab_scores'
    os.makedirs(predict_ab_fd,exist_ok=True)
    # 准备输出文件夹
    os.makedirs(f'{predict_ab_fd}/output',exist_ok=True)
    os.makedirs(f'{predict_ab_fd}/output/pkl_patches',exist_ok=True)
    # 将 patch 处理为模型需要的 pkl 数据
    logger.info('- Convert patches to pkl..')
    cmd = f"""
docker run --rm \\
-e PATH=/home/miniconda3/bin:$PATH \\
-v {ab_struct_predict_fd}/output/ag-TransformedAg/patches/agAbPatch:/home/main/data/test/pdb \\
-v {predict_ab_fd}/output/pkl_patches:/home/main/data/test/pkl \\
-w /home/main/utils \\
{IMAGES['predict_patches']} \\
python getPkl.py \\
    --pdb_path /home/main/data/test/pdb \\
    --pkl /home/main/data/test/pkl
""".strip()
    os.system(cmd)
    # 在上面输出的 pkl 数据上预测，输出每个 patch 文件对应的得分
    logger.info('- Predict by model..')
    if gpu == '-1':
        cmd = f"""
docker run --rm --shm-size 4G \\
-e CUDA_VISIBLE_DEVICES={gpu} \\
-e PATH=/home/miniconda3/bin:$PATH \\
-v {predict_ab_fd}/output/pkl_patches:/home/main/data/test/pkl \\
-v {predict_ab_fd}/output:/home/main/data/test/save \\
-w /home/main/utils \\
{IMAGES['predict_patches']} \\
python predict.py \\
    --data_path /home/main/data/test/pkl/ \\
    --save_path /home/main/data/test/save/patches-scores.pkl
""".strip()
    else:
        cmd = f"""
docker run --rm --gpus all --shm-size 4G \\
-e CUDA_VISIBLE_DEVICES={gpu} \\
-e NVIDIA_DRIVER_CAPABILITIES=compute \\
-e PATH=/home/miniconda3/bin:$PATH \\
-v {predict_ab_fd}/output/pkl_patches:/home/main/data/test/pkl \\
-v {predict_ab_fd}/output:/home/main/data/test/save \\
-w /home/main/utils \\
{IMAGES['predict_patches']} \\
python predict.py \\
    --data_path /home/main/data/test/pkl/ \\
    --save_path /home/main/data/test/save/patches-scores.pkl
""".strip()
    os.system(cmd)
    
    # 4. 将切的 patch 进行可视化
    logger.info('Visualize antigen patches..')
    visualize_patch_fd = f'{tmp_fd}/4-visualize_patch'
    os.makedirs(visualize_patch_fd,exist_ok=True)
    os.makedirs(f'{visualize_patch_fd}/output',exist_ok=True)
    cmd = f"""
docker run --rm \
-v {ab_struct_predict_fd}/input:/job/ag \\
-v {ab_struct_predict_fd}/output/ag-TransformedAg/patches/Agpatch:/job/ag_patch \\
-v {visualize_patch_fd}/output:/job/pse_saved \\
{IMAGES['visualize_patches']} \\
    --ag_path /job/ag/ \\
    --ag_pdbname ag-TransformedAg.pdb \\
    --patch_path /job/ag_patch/ \\
    --pse_savedpath /job/pse_saved/ \\
    --pse_name ag-patches.pse
""".strip()
    os.system(cmd)
    
    # 4. 基于模型预测的 patch 和抗体之间的得分选择指定数量的抗体用于 binning
    logger.info('Select antibodies for binning..')
    binning_fd = f'{tmp_fd}/5-epitope_binning'
    os.makedirs(binning_fd,exist_ok=True)
    os.makedirs(f'{binning_fd}/output',exist_ok=True)
    cmd = f"""
docker run --rm \\
-v {predict_ab_fd}/output:/input/patch_score \\
-v {ab_struct_predict_fd}/output/ag-TransformedAg/Igfold:/input/ab_pdb \\
-v {binning_fd}/output:/output \\
{IMAGES['epitope_binning']} \\
--ppi_score_pkl /input/patch_score/patches-scores.pkl \\
--ab_pdb_fd /input/ab_pdb \\
--save_fd /output \\
-k {' '.join([str(e) for e in k])}
""".strip()
    os.system(cmd)

    # 挑选中间结果保存到最终的结果文件夹中
    logger.info('Saving result..')
    if save_fd.endswith('/'):
        save_fd = save_fd[:-1]
    os.makedirs(save_fd, exist_ok=True)
    os.system(f'cp {visualize_patch_fd}/output/ag-patches.pse {save_fd}/')
    os.system(f'cp -r {binning_fd}/output/* {save_fd}/')
    
    logger.info('Done.')
    

def get_parser():
    parser = argparse.ArgumentParser()
    parser.add_argument('--ag_pdb_fp', type=str, required=True, help='The path of the antigen pdb file.')
    parser.add_argument('--ab_csv_fp', type=str, required=True, help='The path of the antibodies csv file, 3 columns naming name, VH and VL')
    parser.add_argument('-k', type=int, nargs='+', default=[5, 10, 15], help='The number of antibodies to choose for binning, default is [5, 10, 15]')
    parser.add_argument('--gpu', type=str, default='0', help='GPU device to use, deault is 0')
    parser.add_argument('--save_fd', type=str, default='./inspired-result', help='Result of silicon antibodies binning to save, default is ./inspired-result')
    parser.add_argument('--tmp_fd', type=str, default='./.inspired_tmp_data', help='Result of middle data to save, default is ./.inspired_tmp_data')
    args = parser.parse_args()
    return args


if __name__ == '__main__':
    args = get_parser()
    pipeline(args.ag_pdb_fp,
             args.ab_csv_fp,
             args.gpu,
             args.k,
             args.save_fd,
             args.tmp_fd)
